import os, sys
from typing import Dict

parent_path = os.path.abspath(os.path.join(__file__, *([".."] * 4)))
sys.path.insert(0, parent_path)

import traceback
import numpy as np
import socket, json, cv2, yaml, time
from typing import *
from core.task.modules.processors.base_processor import BaseProcessor
from core.utils.PIL_draw import pic_text
from core.utils.visualize import get_color_map_list
from datetime import datetime, timedelta


class VIDEOProcessor(BaseProcessor):
    """直播流检测后处理方案"""

    def __init__(
        self,
        keys,
        Custom_cls_map,
        Custom_target_cls,
        Custom_num_cls,
        Custom_landslide_cls_map,
        Custom_landslide_target_cls,
        Custom_landslide_num_cls,
        Custom_pest_cls_map,
        Custom_pest_target_cls,
        Custom_pest_num_cls,
        VIDEOProcessor_video_writer_path,
        VIDEOProcessor_video_output_directory,
    ) -> None:
        self.keys = keys
        if isinstance(Custom_target_cls, str):
            with open(Custom_target_cls, "r") as f:
                self.target_cls = set(yaml.safe_load(f)["label_list"])
        else:
            self.target_cls = set(Custom_target_cls)
        if isinstance(Custom_cls_map, str):
            with open(Custom_cls_map, "r") as f:
                self.cls_map = {
                    x: y for x, y in enumerate(yaml.safe_load(f)["label_list"])
                }
        else:
            self.cls_map = Custom_cls_map
        if isinstance(Custom_num_cls, str):
            with open(Custom_num_cls, "r") as f:
                self.Custom_num_cls = {
                    x: y for x, y in enumerate(yaml.safe_load(f)["label_list"])
                }
        else:
            self.cls_num = Custom_num_cls
        self.color_list = get_color_map_list(len(self.cls_map))

        if isinstance(Custom_landslide_target_cls, str):
            with open(Custom_landslide_target_cls, "r") as f:
                self.landslide_target_cls = set(yaml.safe_load(f)["label_list"])
        else:
            self.landslide_target_cls = set(Custom_landslide_target_cls)
        if isinstance(Custom_landslide_cls_map, str):
            with open(Custom_landslide_cls_map, "r") as f:
                self.landslide_cls_map = {
                    x: y for x, y in enumerate(yaml.safe_load(f)["label_list"])
                }
        else:
            self.landslide_cls_map = Custom_landslide_cls_map
        if isinstance(Custom_landslide_num_cls, str):
            with open(Custom_landslide_num_cls, "r") as f:
                self.Custom_landslide_num_cls = {
                    x: y for x, y in enumerate(yaml.safe_load(f)["label_list"])
                }
        else:
            self.landslide_cls_num = Custom_landslide_num_cls
        self.landslide_color_list = get_color_map_list(len(self.landslide_cls_map))

        if isinstance(Custom_pest_target_cls, str):
            with open(Custom_pest_target_cls, "r") as f:
                self.pest_target_cls = set(yaml.safe_load(f)["label_list"])
        else:
            self.pest_target_cls = set(Custom_pest_target_cls)
        if isinstance(Custom_pest_cls_map, str):
            with open(Custom_pest_cls_map, "r") as f:
                self.pest_cls_map = {
                    x: y for x, y in enumerate(yaml.safe_load(f)["label_list"])
                }
        else:
            self.pest_cls_map = Custom_pest_cls_map
        if isinstance(Custom_pest_num_cls, str):
            with open(Custom_pest_num_cls, "r") as f:
                self.Custom_pest_num_cls = {
                    x: y for x, y in enumerate(yaml.safe_load(f)["label_list"])
                }
        else:
            self.pest_cls_num = Custom_pest_num_cls
        self.pest_color_list = get_color_map_list(len(self.pest_cls_map))

        self.video_writer_path = VIDEOProcessor_video_writer_path
        self.video_output_directory = VIDEOProcessor_video_output_directory

    def init_check(self):
        assert "det" in self.keys
        assert "detect_landslide" in self.keys
        assert "detect_pest" in self.keys
        assert not self.video_writer_path is None
        assert not self.video_output_directory is None

    def __call__(self, data: Dict) -> Dict:

        batch_img = data[self.keys["in"]][0].copy()
        det_result = data[self.keys["det"]]
        detect_landslide_result = data[self.keys["detect_landslide"]]
        detect_pest_result = data[self.keys["detect_pest"]]

        for img_idx, (bboxes_fire, bboxes_landslide, bboxes_pest) in enumerate(
            zip(
                det_result["boxes"],
                detect_landslide_result["boxes"],
                detect_pest_result["boxes"],
            )
        ):
            draw_thickness = min(batch_img.shape[:2]) // 320
            img_objects = []
            img = batch_img.copy()
            # 林火烟雾模型
            for bbox in bboxes_fire:
                cls_id, score, x1, y1, x2, y2 = bbox
                cls_id = int(cls_id)
                if self.cls_map[cls_id] in self.target_cls:
                    object = {
                        "box": [cls_id, score, x1, y1, x2, y2],
                        "cls_name": self.cls_map[cls_id],
                        "cls_id": self.cls_num[cls_id],
                    }
                    img_objects.append(object)
                    # vis
                    img = cv2.rectangle(
                        img,
                        (int(x1), int(y1)),
                        (int(x2), int(y2)),
                        color=self.color_list[cls_id],
                        thickness=draw_thickness,
                    )
                    text = self.cls_map[cls_id]
                    img = pic_text(
                        img,
                        text,
                        (int(x1), int(y1)),
                        self.color_list[cls_id],
                        draw_thickness * 10,
                    )
            # 滑坡+车辆+行人+生态破坏模型
            for bbox in bboxes_landslide:
                cls_id, score, x1, y1, x2, y2 = bbox
                cls_id = int(cls_id)
                if self.landslide_cls_map[cls_id] in self.landslide_target_cls:
                    object = {
                        "box": [cls_id, score, x1, y1, x2, y2],
                        "cls_name": self.landslide_cls_map[cls_id],
                        "cls_id": self.landslide_cls_num[cls_id],
                    }
                    img_objects.append(object)
                    # vis
                    img = cv2.rectangle(
                        img,
                        (int(x1), int(y1)),
                        (int(x2), int(y2)),
                        color=self.landslide_color_list[cls_id],
                        thickness=draw_thickness,
                    )
                    text = self.landslide_cls_map[cls_id]
                    img = pic_text(
                        img,
                        text,
                        (int(x1), int(y1)),
                        self.landslide_color_list[cls_id],
                        draw_thickness * 10,
                    )
            # 病虫害模型
            for bbox in bboxes_pest:
                cls_id, score, x1, y1, x2, y2 = bbox
                cls_id = int(cls_id)
                if self.pest_cls_map[cls_id] in self.pest_target_cls:
                    object = {
                        "box": [cls_id, score, x1, y1, x2, y2],
                        "cls_name": self.pest_cls_map[cls_id],
                        "cls_id": self.pest_cls_num[cls_id],
                    }
                    img_objects.append(object)
                    # vis
                    img = cv2.rectangle(
                        img,
                        (int(x1), int(y1)),
                        (int(x2), int(y2)),
                        color=self.pest_color_list[cls_id],
                        thickness=draw_thickness,
                    )
                    text = self.pest_cls_map[cls_id]
                    img = pic_text(
                        img,
                        text,
                        (int(x1), int(y1)),
                        self.pest_color_list[cls_id],
                        draw_thickness * 10,
                    )

        # 输出模型检测完数据
        data[self.keys["out_vis"]] = img

        # 视频读取完成,执行 VideoTranscode 和 HttpRequest 模块
        if data.get(self.keys["finish_tag"], False):
            file_name = os.path.basename(data[self.keys["SteamReader_path"]])

            # VideoTranscode
            data[self.keys["video_writer"]] = self.video_writer_path
            data[self.keys["video_output"]] = (
                self.video_output_directory + "vis_" + file_name
            )

            # HttpRequest
            data[self.keys["out_obj"]] = {
                "code": 200,
                # 原视频名称
                "video_path": file_name,
                # 检测后视频名称
                "save_path": "vis_" + file_name,
            }

        return data

    # def close(self):
    #     self.status = 0


if __name__ == "__main__":
    cfg_file = (
        "/home/smartgis/workspace/project/smart_server/projects/smoke/configs/test.yml"
    )
    with open(cfg_file, "r") as f:
        cfg = yaml.safe_load(f)
    # 取对应的内部配置
    cfg = cfg["VIDEOProcessor"]
    video_processor = VIDEOProcessor(**cfg)
